package cn.itcast.demo;

import cn.itcast.demo.pojo.Goods;
import com.alibaba.fastjson.JSON;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpHost;
import org.apache.lucene.search.join.ScoreMode;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.script.mustache.SearchTemplateRequest;
import org.elasticsearch.script.mustache.SearchTemplateResponse;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.nested.Nested;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.search.suggest.Suggest;
import org.elasticsearch.search.suggest.SuggestBuilder;
import org.elasticsearch.search.suggest.SuggestBuilders;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author 虎哥
 */
public class EsDemo {

    private RestHighLevelClient client;

    @Test
    public void testCreateIndexAndMapping() throws IOException {
        // 创建索引库的请求对象，并制定索引库名称
        CreateIndexRequest request = new CreateIndexRequest("goods");
        // settings和mappings
        request.source("{\n" +
                "  \"settings\": {\n" +
                "    \"analysis\": {\n" +
                "      \"analyzer\": {\n" +
                "        \"my_pinyin\": {\n" +
                "          \"tokenizer\": \"ik_smart\",\n" +
                "          \"filter\": [\n" +
                "            \"py\"\n" +
                "          ]\n" +
                "        }\n" +
                "      },\n" +
                "      \"filter\": {\n" +
                "        \"py\": {\n" +
                "          \"type\": \"pinyin\",\n" +
                "          \"keep_full_pinyin\": false,\n" +
                "          \"keep_joined_full_pinyin\": true,\n" +
                "          \"keep_original\": true,\n" +
                "          \"limit_first_letter_length\": 16,\n" +
                "          \"remove_duplicated_term\": true\n" +
                "        }\n" +
                "      }\n" +
                "    }\n" +
                "  },\n" +
                "  \"mappings\": {\n" +
                "    \"properties\": {\n" +
                "      \"id\": {\n" +
                "        \"type\": \"keyword\"\n" +
                "      },\n" +
                "      \"name\": {\n" +
                "        \"type\": \"completion\",\n" +
                "        \"analyzer\": \"my_pinyin\",\n" +
                "        \"search_analyzer\": \"ik_smart\"\n" +
                "      },\n" +
                "      \"title\":{\n" +
                "        \"type\": \"text\",\n" +
                "        \"analyzer\": \"my_pinyin\",\n" +
                "        \"search_analyzer\": \"ik_smart\"\n" +
                "      },\n" +
                "      \"price\":{\n" +
                "        \"type\": \"long\"\n" +
                "      }\n" +
                "    }\n" +
                "  }\n" +
                "}", XContentType.JSON);
        // 创建索引库
        client.indices().create(request, RequestOptions.DEFAULT);
    }

    @Test
    public void testDeleteIndex() throws IOException {
        // 创建一个删除索引的请求对象
        DeleteIndexRequest request = new DeleteIndexRequest("goods");
        // 删除
        client.indices().delete(request, RequestOptions.DEFAULT);
    }

    @Test
    public void testAddDocument() throws IOException {
        Goods goods = new Goods(1L, "红米9", "红米9手机 数码", 1499L);

        // 创建一个新增文档的请求对象
        IndexRequest request = new IndexRequest("goods")
                .id(goods.getId().toString())
                .source(JSON.toJSONString(goods), XContentType.JSON);

        // 添加一条文档到索引库
        IndexResponse response = client.index(request, RequestOptions.DEFAULT);

        System.out.println("response = " + response);
    }

    @Test
    public void testBulk() throws IOException {
        // 1.准备文档数据
        List<Goods> list = new ArrayList<>();
        list.add(new Goods(1L, "红米9", "红米9手机 数码", 1499L));
        list.add(new Goods(2L, "三星 Galaxy A90", "三星 Galaxy A90 手机 数码 疾速5G 骁龙855", 3099L));
        list.add(new Goods(3L, "Sony WH-1000XM3", "Sony WH-1000XM3 降噪耳机 数码", 2299L));
        list.add(new Goods(4L, "松下剃须刀", "松下电动剃须刀高转速磁悬浮马达", 599L));

        // 创建一个批处理的请求
        BulkRequest request = new BulkRequest();

        for (Goods goods : list) {
            request.add(new IndexRequest("goods")
                    .id(goods.getId().toString())
                    .source(JSON.toJSONString(goods), XContentType.JSON));
        }

        // 批处理
        BulkResponse response = client.bulk(request, RequestOptions.DEFAULT);
        System.out.println("response = " + response);
    }

    @Test
    public void testGetDocumentByIdSync() throws IOException {
        // 准备一个查询文档的请求
        GetRequest request = new GetRequest("goods", "1");
        // 查询一个文档，耗时50ms
        GetResponse response = client.get(request, RequestOptions.DEFAULT);

        System.out.println("response = " + response);

        // 获取source
        String json = response.getSourceAsString();
        // 把json反序列化
        Goods goods = JSON.parseObject(json, Goods.class);
        System.out.println("goods = " + goods);
    }

    @Test
    public void testGetDocumentByIdAsync() throws IOException, InterruptedException {

        System.out.println("准备开始查询");
        // 准备一个查询文档的请求
        GetRequest request = new GetRequest("goods", "1");
        // 异步查询一个文档，耗时50ms
        client.getAsync(request, RequestOptions.DEFAULT, new ActionListener<GetResponse>() {
            @Override
            public void onResponse(GetResponse response) {
                // 获取source
                String json = response.getSourceAsString();
                // 把json反序列化
                Goods goods = JSON.parseObject(json, Goods.class);

                System.out.println("查询结束，得到结果： " + goods);
            }

            @Override
            public void onFailure(Exception e) {
                e.printStackTrace();
            }
        });

        System.out.println("请求已经发出，等待执行结果！");


        Thread.sleep(2000);
    }

    @Test
    public void testMono() throws InterruptedException {
        Mono<Goods> mono = Mono.create(monoSink -> {
            // 准备一个查询文档的请求
            GetRequest request = new GetRequest("goods", "1");
            // 异步查询一个文档，耗时50ms
            client.getAsync(request, RequestOptions.DEFAULT, new ActionListener<GetResponse>() {
                @Override
                public void onResponse(GetResponse response) {
                    // 获取source
                    String json = response.getSourceAsString();
                    // 把json反序列化
                    Goods goods = JSON.parseObject(json, Goods.class);
                    monoSink.success(goods);
                }

                @Override
                public void onFailure(Exception e) {
                    monoSink.error(e);
                }
            });
        });

        mono.subscribe(System.out::println);

        Thread.sleep(2000);
    }

    @Test
    public void testDeleteById() throws IOException {
        // 创建一个删除文档的请求对象
        DeleteRequest request = new DeleteRequest("goods", "1");
        // 删除
        DeleteResponse response = client.delete(request, RequestOptions.DEFAULT);
        System.out.println("response = " + response);
    }

    @Test
    public void testQuery() throws IOException {
        // 搜索条件的构建器
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        // 查询条件
        sourceBuilder.query(QueryBuilders.matchQuery("title", "红米手机"));
        // 分页条件
        sourceBuilder.from(0);
        sourceBuilder.size(20);
        // 排序条件
        sourceBuilder.sort("price", SortOrder.DESC);
        // TODO 高亮
        sourceBuilder.highlighter(new HighlightBuilder().field("title"));
        // TODO 聚合

        // 准备一个搜索的请求对象
        SearchRequest request = new SearchRequest("goods");
        // 添加查询条件
        request.source(sourceBuilder);

        // 搜索数据
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);

        SearchHits searchHits = response.getHits();

        SearchHit[] hits = searchHits.getHits();

        for (SearchHit hit : hits) {
            String json = hit.getSourceAsString();
            // 把json反序列化
            Goods goods = JSON.parseObject(json, Goods.class);

            // 获取高亮
            Map<String, HighlightField> map = hit.getHighlightFields();
            HighlightField field = map.get("title");
            String value = StringUtils.join(field.getFragments());
            goods.setTitle(value);

            System.out.println("goods = " + goods);
        }
    }

    @Test
    public void testNestedQuery() throws IOException {
        // 搜索条件的构建器
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        // 查询条件
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
        queryBuilder.must(QueryBuilders.matchQuery("user.first", "John"));
        queryBuilder.must(QueryBuilders.matchQuery("user.last", "White"));

        sourceBuilder.query(QueryBuilders.nestedQuery("user", queryBuilder, ScoreMode.None));
        // 准备一个搜索的请求对象
        SearchRequest request = new SearchRequest("my_index");
        // 添加查询条件
        request.source(sourceBuilder);
        // 搜索数据
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);

        SearchHits hits = response.getHits();

        for (SearchHit hit : hits.getHits()) {
            System.out.println("hit = " + hit.getSourceAsString());
        }
    }

    @Test
    public void testNestedAggregations() throws IOException {
        // 搜索条件的构建器
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        // size置为0，不要hits数据
        sourceBuilder.size(0);
        // 聚合
        sourceBuilder.aggregation(
                AggregationBuilders.nested("nestedAgg", "user")
                        .subAggregation(AggregationBuilders.terms("firstAgg").field("user.first")
                                .subAggregation(AggregationBuilders.terms("lastAgg").field("user.last")))
        );


        // 准备一个搜索的请求对象
        SearchRequest request = new SearchRequest("my_index");
        // 添加查询条件
        request.source(sourceBuilder);

        // 搜索数据
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);

        // 获取聚合结果
        Aggregations aggregations = response.getAggregations();

        // 获取了nested的聚合结果
        Nested nestedAgg = aggregations.get("nestedAgg");

        // 获取第一层聚合
        Terms firstAgg = nestedAgg.getAggregations().get("firstAgg");

        for (Terms.Bucket bucket : firstAgg.getBuckets()) {
            System.out.println("firstName: " + bucket.getKeyAsString() + ", 包含下列lastName：");
            // 获取第二层聚合
            Terms lastAgg = bucket.getAggregations().get("lastAgg");
            for (Terms.Bucket lastAggBucket : lastAgg.getBuckets()) {
                System.out.println("\t lastName: " + lastAggBucket.getKeyAsString());
            }
        }
    }

    @Test
    public void testSuggest() throws IOException {
        // 搜索条件的构建器
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        // 添加自动补全功能
        sourceBuilder.suggest(new SuggestBuilder().addSuggestion(
                "nameSuggest",
                SuggestBuilders.completionSuggestion("name").prefix("s"))
        );
        // 准备一个搜索的请求对象
        SearchRequest request = new SearchRequest("goods");
        // 添加查询条件
        request.source(sourceBuilder);
        // 搜索
        SearchResponse response = client.search(request, RequestOptions.DEFAULT);

        // 获取补全结果
        Suggest suggest = response.getSuggest();

        suggest.getSuggestion("nameSuggest")
                .forEach(os -> os.getOptions().forEach(o -> System.out.println((o.getText()))));
    }

    @Test
    public void testQueryTemplate() throws IOException {
        // 创建模板查询的请求对象
        SearchTemplateRequest request = new SearchTemplateRequest(new SearchRequest("my_index"));
        // 指定模板id
        request.setScript("aggTemplate");
        // 指定模板类型
        request.setScriptType(ScriptType.STORED);
        // 指定模板参数
        Map<String, Object> params = new HashMap<>();
        params.put("path", "user");
        params.put("field1", "user.first");
        params.put("field2", "user.last");
        request.setScriptParams(params);


        SearchTemplateResponse templateResponse = client.searchTemplate(request, RequestOptions.DEFAULT);
        // 搜索数据
        SearchResponse response = templateResponse.getResponse();

        // 获取聚合结果
        Aggregations aggregations = response.getAggregations();

        // 获取了nested的聚合结果
        Nested nestedAgg = aggregations.get("nestedAgg");

        // 获取第一层聚合
        Terms firstAgg = nestedAgg.getAggregations().get("firstAgg");

        for (Terms.Bucket bucket : firstAgg.getBuckets()) {
            System.out.println("firstName: " + bucket.getKeyAsString() + ", 包含下列lastName：");
            // 获取第二层聚合
            Terms lastAgg = bucket.getAggregations().get("lastAgg");
            for (Terms.Bucket lastAggBucket : lastAgg.getBuckets()) {
                System.out.println("\t lastName: " + lastAggBucket.getKeyAsString());
            }
        }
    }

    /**
     * 建立连接
     */
    @Before
    public void init() throws IOException {
        client = new RestHighLevelClient(
                RestClient.builder(
                        new HttpHost("192.168.150.101", 9200, "http")
                )
        );
    }

    /**
     * 关闭客户端连接
     */
    @After
    public void close() throws IOException {
        client.close();
    }
}
